import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.function.VoidFunction;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class HdfsSparkSync {
    
    private final FileSystem srcFs;  // 源集群FileSystem
    private final FileSystem dstFs;  // 目标集群FileSystem
    private final long bandwidthLimit; // 带宽限制，单位：字节/秒

    public HdfsSparkSync(Configuration srcConf, Configuration dstConf, long bandwidthLimit) throws IOException {
        // 初始化源集群FileSystem
        srcFs = FileSystem.get(srcConf);
        dstFs = FileSystem.get(dstConf);
        this.bandwidthLimit = bandwidthLimit * 1024 * 1024; // 将带宽转换为字节/秒
    }

    /**
     * update - 在Driver上列出源和目标文件，分发需要复制的文件列表给Executors并复制
     */
    public void update(Path srcDir, Path dstDir, JavaSparkContext sc) throws IOException {
        List<Path> filesToCopy = new ArrayList<>();
        gatherFilesToUpdate(srcDir, dstDir, filesToCopy);
        
        // 将文件路径列表转化为RDD并分发到Executor
        JavaRDD<Path> filesRDD = sc.parallelize(filesToCopy);
        filesRDD.foreach((VoidFunction<Path>) path -> {
            Configuration srcConf = new Configuration();
            srcConf.set("hadoop.security.authentication", "Kerberos");
            FileSystem srcFs = FileSystem.get(srcConf);
            FileSystem dstFs = FileSystem.get(new Configuration());
            
            UserGroupInformation.setConfiguration(srcConf);
            UserGroupInformation.loginUserFromKeytab("your_principal@YOUR_REALM", "/path/to/your.keytab");

            Path dstPath = new Path(dstDir, path.getName());
            copyFile(srcFs, dstFs, path, dstPath);
            preservePermissions(srcFs, dstFs, path, dstPath);
        });
    }

    /**
     * delete - Driver上遍历目标集群文件并删除源集群上不存在的文件
     */
    public void delete(Path srcDir, Path dstDir) throws IOException {
        FileStatus[] dstFiles = dstFs.listStatus(dstDir);

        for (FileStatus dstFile : dstFiles) {
            Path dstPath = dstFile.getPath();
            Path srcPath = new Path(srcDir, dstPath.getName());

            if (dstFile.isDirectory()) {
                delete(srcPath, dstPath);
                if (dstFs.listStatus(dstPath).length == 0) {
                    dstFs.delete(dstPath, true);
                    System.out.println("Deleted empty directory: " + dstPath);
                }
            } else {
                if (!srcFs.exists(srcPath)) {
                    dstFs.delete(dstPath, true);
                    System.out.println("Deleted: " + dstPath);
                }
            }
        }
    }

    /**
     * gatherFilesToUpdate - 递归遍历源目录，记录需要更新的文件
     */
    private void gatherFilesToUpdate(Path srcDir, Path dstDir, List<Path> filesToCopy) throws IOException {
        FileStatus[] srcFiles = srcFs.listStatus(srcDir);

        if (!dstFs.exists(dstDir)) {
            dstFs.mkdirs(dstDir);
        }

        for (FileStatus srcFile : srcFiles) {
            Path srcPath = srcFile.getPath();
            Path dstPath = new Path(dstDir, srcPath.getName());

            if (srcFile.isDirectory()) {
                gatherFilesToUpdate(srcPath, dstPath, filesToCopy);
            } else {
                if (!dstFs.exists(dstPath) || srcFs.getFileStatus(srcPath).getModificationTime() > dstFs.getFileStatus(dstPath).getModificationTime()) {
                    filesToCopy.add(srcPath);
                }
            }
        }
    }

    /**
     * copyFile - 复制文件并限制带宽
     */
    private void copyFile(FileSystem srcFs, FileSystem dstFs, Path srcPath, Path dstPath) throws IOException {
        try (var inStream = srcFs.open(srcPath);
             var outStream = dstFs.create(dstPath, true)) {

            byte[] buffer = new byte[4096];
            long startTime = System.currentTimeMillis();
            int bytesRead;
            long totalBytesRead = 0;

            while ((bytesRead = inStream.read(buffer)) > 0) {
                outStream.write(buffer, 0, bytesRead);
                totalBytesRead += bytesRead;

                if (totalBytesRead >= bandwidthLimit) {
                    long elapsedTime = System.currentTimeMillis() - startTime;
                    if (elapsedTime < 1000) {
                        Thread.sleep(1000 - elapsedTime);
                    }
                    totalBytesRead = 0;
                    startTime = System.currentTimeMillis();
                }
            }
        }
        System.out.println("Copied: " + srcPath + " to " + dstPath);
    }

    /**
     * preservePermissions - 保留文件权限
     */
    private void preservePermissions(FileSystem srcFs, FileSystem dstFs, Path srcPath, Path dstPath) throws IOException {
        FileStatus srcStatus = srcFs.getFileStatus(srcPath);
        dstFs.setPermission(dstPath, srcStatus.getPermission());
        dstFs.setOwner(dstPath, srcStatus.getOwner(), srcStatus.getGroup());
    }

    public static void main(String[] args) throws Exception {
        Configuration srcConf = new Configuration();
        Configuration dstConf = new Configuration();

        srcConf.set("fs.defaultFS", "hdfs://source_cluster");
        dstConf.set("fs.defaultFS", "hdfs://destination_cluster");

        UserGroupInformation.setConfiguration(srcConf);
        UserGroupInformation.loginUserFromKeytab("your_principal@YOUR_REALM", "/path/to/your.keytab");

        JavaSparkContext sc = new JavaSparkContext();
        HdfsSparkSync sync = new HdfsSparkSync(srcConf, dstConf, 20); // 设置带宽限制为20 MB/s
        Path srcDir = new Path("/source/path");
        Path dstDir = new Path("/destination/path");

        sync.delete(srcDir, dstDir);  // 删除多余文件
        sync.update(srcDir, dstDir, sc);  // 更新文件
        sc.close();
    }
}
